from transformers import *

import tensorflow as tf
import numpy as np
from tqdm import tqdm
from encoder import BaseEncoder
import re


class RegexEncoder(BaseEncoder):

    def __init__(self, regex_list, max_len):
        self.max_len = max_len

        self.regex_list = regex_list



    def regex_encoder(self,text):
        embed = np.zeros((self.max_len, len(self.regex_list)))

        for i, regex_str in enumerate(self.regex_list):

            for item in re.finditer(regex_str, text):

                embed[item.start():item.end(), i] = 1

        return embed
    
    def build(self):

        input_ids_l = tf.keras.layers.Input(
            (self.max_len,len(self.regex_list)), dtype=tf.float32, name=f"regex_input_ids")
        encoder = tf.keras.layers.BatchNormalization()(input_ids_l)

        encoder = tf.keras.layers.Bidirectional(
            tf.keras.layers.GRU(128,return_sequences=True))(encoder)

        encoder = tf.keras.layers.LeakyReLU()(encoder)

        return input_ids_l,encoder
    


    def preprocess_input(self, text_list):


        embed_list = []

        for text in tqdm(text_list):

            embed_list.append(self.regex_encoder(text))

        return np.array(embed_list)
